// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/prestosql/types/ip_prefix_registration.h>

#include <pollux/expression/cast_expr.h>
#include <pollux/functions/prestosql/types/ip_prefix_type.h>

namespace kumo::pollux {
    namespace {
        class IPPrefixCastOperator : public exec::CastOperator {
        public:
            bool isSupportedFromType(const TypePtr &other) const override {
                switch (other->kind()) {
                    case TypeKind::VARCHAR:
                        return true;
                    case TypeKind::HUGEINT:
                        return isIPAddressType(other);
                    default:
                        return false;
                }
            }

            bool isSupportedToType(const TypePtr &other) const override {
                switch (other->kind()) {
                    case TypeKind::VARCHAR:
                        return true;
                    case TypeKind::HUGEINT:
                        return isIPAddressType(other);
                    default:
                        return false;
                }
            }

            void castTo(
                const BaseVector &input,
                exec::EvalCtx &context,
                const SelectivityVector &rows,
                const TypePtr &resultType,
                VectorPtr &result) const override {
                context.ensure_writable(rows, resultType, result);
                switch (input.type_kind()) {
                    case TypeKind::VARCHAR:
                        return castFromString(input, context, rows, *result);
                    case TypeKind::HUGEINT: {
                        if (isIPAddressType(input.type())) {
                            return castFromIpAddress(input, context, rows, *result);
                        }
                        [[fallthrough]];
                    }
                    default:
                        POLLUX_NYI(
                            "Cast from {} to IPPrefix not yet supported",
                            input.type()->toString());
                }
            }

            void castFrom(
                const BaseVector &input,
                exec::EvalCtx &context,
                const SelectivityVector &rows,
                const TypePtr &resultType,
                VectorPtr &result) const override {
                context.ensure_writable(rows, resultType, result);
                switch (resultType->kind()) {
                    case TypeKind::VARCHAR:
                        return castToString(input, context, rows, *result);
                    default:
                        POLLUX_NYI(
                            "Cast from IPPrefix to {} not yet supported",
                            resultType->toString());
                }
            }

        private:
            static void castToString(
                const BaseVector &input,
                exec::EvalCtx &context,
                const SelectivityVector &rows,
                BaseVector &result) {
                auto *flatResult = result.as<FlatVector<StringView> >();
                auto row_vector = input.as<RowVector>();
                const auto *ipaddr = row_vector->childAt(ipaddress::kIpRowIndex)
                        ->as<SimpleVector<int128_t> >();
                const auto *prefix = row_vector->childAt(ipaddress::kIpPrefixRowIndex)
                        ->as<SimpleVector<int8_t> >();
                context.applyToSelectedNoThrow(rows, [&](auto row) {
                    const auto ipAddrVal = ipaddr->value_at(row);
                    // The string representation of the last byte needs
                    // to be unsigned
                    const uint8_t prefixVal = prefix->value_at(row);

                    // Copy the first 16 bytes into a ByteArray16.
                    melon::ByteArray16 addrBytes;
                    memcpy(&addrBytes, &ipAddrVal, ipaddress::kIPAddressBytes);
                    // Reverse the bytes to get the correct order. Similar to
                    // IPAddressType. We assume we're ALWAYS on a little endian machine.
                    // Note: for big endian, we should not reverse the bytes.
                    std::reverse(addrBytes.begin(), addrBytes.end());
                    // // Construct a V6 address from the ByteArray16.
                    melon::IPAddressV6 v6Addr(addrBytes);

                    // Inline func to get string for ipv4 or ipv6 string
                    const auto ipString =
                            (v6Addr.isIPv4Mapped()) ? v6Addr.createIPv4().str() : v6Addr.str();

                    // Format of string is {ipString}/{mask}
                    auto stringRet = fmt::format("{}/{}", ipString, prefixVal);

                    // Write the string to the result vector
                    exec::StringWriter result(flatResult, row);
                    result.append(stringRet);
                    result.finalize();
                });
            }

            static void castFromIpAddress(
                const BaseVector &input,
                exec::EvalCtx &context,
                const SelectivityVector &rows,
                BaseVector &result) {
                auto *rowVectorResult = result.as<RowVector>();
                auto intIpAddrVec =
                        rowVectorResult->childAt(0)->as_checked<FlatVector<int128_t> >();
                auto intPrefixVec =
                        rowVectorResult->childAt(1)->as_checked<FlatVector<int8_t> >();
                DecodedVector decoded(input, rows);

                context.applyToSelectedNoThrow(rows, [&](auto row) {
                    auto intIpAddr = decoded.value_at<int128_t>(row);
                    const auto tryPrefixLength =
                            ipaddress::tryIpPrefixLengthFromIPAddressType(intIpAddr);
                    if (MELON_UNLIKELY(tryPrefixLength.hasError())) {
                        context.setStatus(row, std::move(tryPrefixLength.error()));
                        return;
                    }
                    intIpAddrVec->set(row, intIpAddr);
                    intPrefixVec->set(row, tryPrefixLength.value());
                });
            }

            static void castFromString(
                const BaseVector &input,
                exec::EvalCtx &context,
                const SelectivityVector &rows,
                BaseVector &result) {
                auto *rowVectorResult = result.as<RowVector>();
                auto intIpAddrVec =
                        rowVectorResult->childAt(0)->as_checked<FlatVector<int128_t> >();
                auto intPrefixVec =
                        rowVectorResult->childAt(1)->as_checked<FlatVector<int8_t> >();

                DecodedVector decoded(input, rows);

                context.applyToSelectedNoThrow(rows, [&](auto row) {
                    auto ipAddressStringView = decoded.value_at<StringView>(row);
                    auto tryIpPrefix = ipaddress::tryParseIpPrefixString(ipAddressStringView);
                    if (tryIpPrefix.hasError()) {
                        context.setStatus(row, std::move(tryIpPrefix.error()));
                        return;
                    }

                    const auto &ipPrefix = tryIpPrefix.value();
                    intIpAddrVec->set(row, ipPrefix.first);
                    intPrefixVec->set(row, ipPrefix.second);
                });
            }
        };

        class IPPrefixTypeFactories : public CustomTypeFactories {
        public:
            TypePtr getType(const std::vector<TypeParameter> &parameters) const override {
                POLLUX_CHECK(parameters.empty());
                return IPPrefixType::get();
            }

            exec::CastOperatorPtr getCastOperator() const override {
                return std::make_shared<IPPrefixCastOperator>();
            }

            AbstractInputGeneratorPtr getInputGenerator(
                const InputGeneratorConfig & /*config*/) const override {
                return nullptr;
            }
        };
    } // namespace

    void registerIPPrefixType() {
        registerCustomType(
            "ipprefix", std::make_unique<const IPPrefixTypeFactories>());
    }
} // namespace kumo::pollux
